package com.castle.fortress.admin.ai.controller;

import cn.hutool.socket.aio.AioClient;
import com.alibaba.druid.support.spring.stat.annotation.Stat;
import com.baomidou.mybatisplus.core.metadata.IPage;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.castle.fortress.admin.ai.config.QwenLlm;
import com.castle.fortress.admin.ai.config.QwenLlmConfig;
import com.castle.fortress.admin.ai.llm.Llm;
import com.castle.fortress.admin.ai.message.AiMessage;
import com.castle.fortress.admin.ai.message.HumanMessage;
import com.castle.fortress.admin.ai.message.MessageStatus;
import com.castle.fortress.admin.ai.prompt.HistoriesPrompt;
import com.castle.fortress.admin.core.constants.GlobalConstants;
import com.castle.fortress.admin.ai.entity.AiRecordEntity;
import com.castle.fortress.admin.ai.dto.AiRecordDto;
import com.castle.fortress.admin.ai.service.AiRecordService;
import com.castle.fortress.common.utils.ConvertUtil;
import com.castle.fortress.common.entity.RespBody;
import com.castle.fortress.common.exception.BizException;
import com.castle.fortress.common.respcode.GlobalRespCode;
import com.castle.fortress.common.enums.OperationTypeEnum;
import com.castle.fortress.admin.core.annotation.CastleLog;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import springfox.documentation.annotations.ApiIgnore;
import org.apache.logging.log4j.util.Strings;
import org.apache.shiro.authz.annotation.RequiresPermissions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.stereotype.Controller;
import com.castle.fortress.admin.utils.ExcelUtils;

import javax.servlet.http.HttpServletResponse;

import com.castle.fortress.common.entity.DynamicExcelEntity;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;

/**
 * ai信息记录表 控制器
 *
 * @author
 * @since 2024-03-09
 */
@Api(tags = "ai信息记录表管理控制器")
@Controller
public class AiRecordController {
    @Autowired
    private AiRecordService aiRecordService;
    private static HashMap<String, HistoriesPrompt> promptHashMap = new HashMap<String, HistoriesPrompt>();
    private static HashMap<String, Llm> llmHashMap = new HashMap<String, Llm>();
    private HashMap<String, Object> hisData = new HashMap<String, Object>();

//    private static QwenLlmConfig config;
//
//    static {
//        config = new QwenLlmConfig();
//        config.setApiKey("sk-887ef1c8383e4027b014f85372107fda");
//        config.setModel("qwen-max");
//    }

    /**
     * ai信息记录表的分页展示
     *
     * @param aiRecordDto ai信息记录表实体类
     * @param current     当前页
     * @param size        每页记录数
     * @return
     */
    @CastleLog(operLocation = "ai信息记录表", operType = OperationTypeEnum.QUERY)
    @ApiOperation("ai信息记录表分页展示")
    @GetMapping("/ai/aiRecord/page")
    @ResponseBody
//    @RequiresPermissions("ai:aiRecord:pageList")
    public RespBody<IPage<AiRecordDto>> pageAiRecord(AiRecordDto aiRecordDto, @RequestParam(required = false) Integer current, @RequestParam(required = false) Integer size) {
        Integer pageIndex = current == null ? GlobalConstants.DEFAULT_PAGE_INDEX : current;
        Integer pageSize = size == null ? GlobalConstants.DEFAULT_PAGE_SIZE : size;
        Page<AiRecordDto> page = new Page(pageIndex, pageSize);
        IPage<AiRecordDto> pages = aiRecordService.pageAiRecord(page, aiRecordDto);

        return RespBody.data(pages);
    }

    /**
     * ai信息记录表的列表展示
     *
     * @param aiRecordDto ai信息记录表实体类
     * @return
     */
    @CastleLog(operLocation = "ai信息记录表", operType = OperationTypeEnum.QUERY)
    @ApiOperation("ai信息记录表列表展示")
    @GetMapping("/ai/aiRecord/list")
    @ResponseBody
    public RespBody<List<AiRecordDto>> listAiRecord(AiRecordDto aiRecordDto) {
        List<AiRecordDto> list = aiRecordService.listAiRecord(aiRecordDto);
        return RespBody.data(list);
    }

    /**
     * ai信息记录表保存
     *
     * @param aiRecordDto ai信息记录表实体类
     * @return
     */
    @CastleLog(operLocation = "ai信息记录表", operType = OperationTypeEnum.INSERT)
    @ApiOperation("ai信息记录表保存")
    @PostMapping("/ai/aiRecord/save")
    @ResponseBody
    @RequiresPermissions("ai:aiRecord:save")
    public RespBody<String> saveAiRecord(@RequestBody AiRecordDto aiRecordDto) {
        if (aiRecordDto == null) {
            throw new BizException(GlobalRespCode.PARAM_MISSED);
        }
        AiRecordEntity aiRecordEntity = ConvertUtil.transformObj(aiRecordDto, AiRecordEntity.class);
        if (aiRecordService.save(aiRecordEntity)) {
            return RespBody.data("保存成功");
        } else {
            return RespBody.fail(GlobalRespCode.OPERATE_ERROR);
        }
    }

    /**
     * ai信息记录表编辑
     *
     * @param aiRecordDto ai信息记录表实体类
     * @return
     */
    @CastleLog(operLocation = "ai信息记录表", operType = OperationTypeEnum.UPDATE)
    @ApiOperation("ai信息记录表编辑")
    @PostMapping("/ai/aiRecord/edit")
    @ResponseBody
    @RequiresPermissions("ai:aiRecord:edit")
    public RespBody<String> updateAiRecord(@RequestBody AiRecordDto aiRecordDto) {
        if (aiRecordDto == null || aiRecordDto.getId() == null || aiRecordDto.getId().equals(0L)) {
            throw new BizException(GlobalRespCode.PARAM_MISSED);
        }
        AiRecordEntity aiRecordEntity = ConvertUtil.transformObj(aiRecordDto, AiRecordEntity.class);
        if (aiRecordService.updateById(aiRecordEntity)) {
            return RespBody.data("保存成功");
        } else {
            return RespBody.fail(GlobalRespCode.OPERATE_ERROR);
        }
    }

    @CastleLog(operLocation = "ai信息记录表", operType = OperationTypeEnum.QUERY)
    @ApiOperation("ai信息记录表编辑")
    @PostMapping("/ai/aiRecord/sendAi")
    @ResponseBody
    public RespBody<String> sendAi(@RequestBody Map<String, String> data) {
        if (data == null || !data.containsKey("groupId") || !data.containsKey("content") || !data.containsKey("userId")) {
            throw new BizException(GlobalRespCode.PARAM_MISSED);
        }
        String key = data.get("userId") + "_" + data.get("groupId");
        if (!llmHashMap.containsKey(key)) {
            QwenLlmConfig config = new QwenLlmConfig();
            config.setApiKey("sk-887ef1c8383e4027b014f85372107fda");
            config.setModel("qwen-max");
            Llm llm = new QwenLlm(config);
            HistoriesPrompt prompt = new HistoriesPrompt();
            List<AiRecordEntity> list = aiRecordService.getByTopN(data.get("groupId"), data.get("userId"));
            for (AiRecordEntity aiRecordEntity : list) {
                prompt.addMessage(new HumanMessage(aiRecordEntity.getUserContent()));
                AiMessage aiMessage = new AiMessage();
                aiMessage.setStatus(MessageStatus.END);
                aiMessage.setTotalTokens(aiRecordEntity.getReturnContent().length());
                aiMessage.setFullContent(aiRecordEntity.getReturnContent());
                prompt.addMessage(aiMessage);
            }
            llmHashMap.put(key, llm);
            promptHashMap.put(key, prompt);
        }
        HistoriesPrompt prompt = promptHashMap.get(key);
        Llm llm = llmHashMap.get(key);
        String id = UUID.randomUUID().toString();
        String content = String.valueOf(data.get("content"));
        System.out.println(content);
        prompt.addMessage(new HumanMessage(content));
        StringBuffer buffer = new StringBuffer();
        HashMap<String, String> map = new HashMap<>();
        llm.chatAsync(prompt, (context, response) -> {
            AiMessage message = response.getMessage();
            buffer.append(message.getContent());
            map.put("data",buffer.toString());
            map.put("status",message.getStatus().name());
            hisData.put(id,map);
            System.out.println(">>>> " + message.getContent());
            if (message.getStatus().name().equals("END")) {
                AiRecordEntity aiRecordEntity=new AiRecordEntity();
                aiRecordEntity.setUserId(Long.parseLong(data.get("userId")));
                aiRecordEntity.setGroupId(Long.parseLong(data.get("groupId")));
                aiRecordEntity.setUserContent(data.get("content"));
                aiRecordEntity.setReturnContent(buffer.toString());
                aiRecordService.save(aiRecordEntity);
            }
        });
        return RespBody.data(id);
    }
    @GetMapping("/ai/aiRecord/getAiData")
    @ResponseBody
    public RespBody<Object> getAiData(@RequestParam String id){
        return RespBody.data(hisData.get(id));
    }


    /**
     * ai信息记录表删除
     *
     * @param id
     * @return
     */
    @CastleLog(operLocation = "ai信息记录表", operType = OperationTypeEnum.DELETE)
    @ApiOperation("ai信息记录表删除")
    @PostMapping("/ai/aiRecord/delete")
    @ResponseBody
    @RequiresPermissions("ai:aiRecord:delete")
    public RespBody<String> deleteAiRecord(@RequestParam Long id) {
        if (id == null) {
            throw new BizException(GlobalRespCode.PARAM_MISSED);
        }
        if (aiRecordService.removeById(id)) {
            return RespBody.data("保存成功");
        } else {
            return RespBody.fail(GlobalRespCode.OPERATE_ERROR);
        }
    }


    /**
     * ai信息记录表批量删除
     *
     * @param ids
     * @return
     */
    @CastleLog(operLocation = "ai信息记录表", operType = OperationTypeEnum.DELETE)
    @ApiOperation("ai信息记录表批量删除")
    @PostMapping("/ai/aiRecord/deleteBatch")
    @ResponseBody
    @RequiresPermissions("ai:aiRecord:deleteBatch")
    public RespBody<String> deleteAiRecordBatch(@RequestBody List<Long> ids) {
        if (ids == null || ids.size() < 1) {
            throw new BizException(GlobalRespCode.PARAM_MISSED);
        }
        if (aiRecordService.removeByIds(ids)) {
            return RespBody.data("保存成功");
        } else {
            return RespBody.fail(GlobalRespCode.OPERATE_ERROR);
        }
    }

    /**
     * ai信息记录表详情
     *
     * @param id ai信息记录表id
     * @return
     */
    @CastleLog(operLocation = "ai信息记录表", operType = OperationTypeEnum.QUERY)
    @ApiOperation("ai信息记录表详情")
    @GetMapping("/ai/aiRecord/info")
    @ResponseBody
    @RequiresPermissions("ai:aiRecord:info")
    public RespBody<AiRecordDto> infoAiRecord(@RequestParam Long id) {
        if (id == null) {
            throw new BizException(GlobalRespCode.PARAM_MISSED);
        }
        AiRecordEntity aiRecordEntity = aiRecordService.getById(id);
        AiRecordDto aiRecordDto = ConvertUtil.transformObj(aiRecordEntity, AiRecordDto.class);

        return RespBody.data(aiRecordDto);
    }

    /**
     * 动态表头导出 依据展示的字段导出对应报表
     *
     * @param dynamicExcelEntity
     * @param response
     * @throws Exception
     */
    @CastleLog(operLocation = "ai信息记录表", operType = OperationTypeEnum.EXPORT)
    @PostMapping("/ai/aiRecord/exportDynamic")
    @ApiOperation("动态表头导出，依据展示的字段导出对应报表")
    public void exportDynamic(@RequestBody DynamicExcelEntity<AiRecordDto> dynamicExcelEntity, HttpServletResponse response) throws Exception {
        List<AiRecordDto> list = aiRecordService.listAiRecord(dynamicExcelEntity.getDto());
        //字典、枚举、接口、json常量等转义后的列表数据 根据实际情况初始化该对象，null为list数据直接导出
        List<List<Object>> dataList = null;
        /**
         * 根据实际情况初始化dataList,可参照 com.castle.fortress.admin.system.controller.TmpDemoController类中方法：exportDynamic
         */
        ExcelUtils.exportDynamic(response, dynamicExcelEntity.getFileName() + ".xlsx", null, list, dynamicExcelEntity.getHeaderList(), dataList);
    }


}
